package udpblock

import (
	"encoding/binary"
	"errors"
	"io"
	"os"
	"sync/atomic"
	"time"
)

// AddLenHead 是在数据的头部增加长度数据，用以在流式输出的连接中传输不粘包的数据
func AddLenHead(data []byte) []byte {
	dataLen := len(data)
	res := make([]byte, dataLen+3)
	copy(res[binary.MaxVarintLen16:], data)
	binary.PutVarint(res[:binary.MaxVarintLen16], int64(dataLen))
	return res
}

// GetData 是从数据中获取真正的数据分片，如果数据并没有使用 AddLenHead 操作过，则该函数不能正确工作
func GetData(data []byte) ([]byte, error) {
	dataLen, i := binary.Varint(data[:binary.MaxVarintLen16])
	if i <= 0 {
		return nil, errors.New("head is error")
	}
	res := make([]byte, dataLen)
	headLen := binary.MaxVarintLen16
	copy(res, data[headLen:headLen+int(dataLen)])
	return res, nil
}

// DeadlineReader 是 io.Reader 加上 SetReadDeadline 方法的接口
type DeadlineReader interface {
	Read(b []byte) (n int, err error)
	SetReadDeadline(t time.Time) error
}

// dataReader 会不断地从 conn 中读取数据，并解析头部的数据，
// 返回数据部分
type dataReader struct {
	conn        DeadlineReader
	dataChannel chan []byte
	errChan     chan error
	closed      atomic.Bool
}

func NewDataReader(conn DeadlineReader) (io.ReadCloser, error) {
	if conn == nil {
		return nil, errors.New("conn is nil")
	}
	res := &dataReader{
		conn:        conn,
		dataChannel: make(chan []byte, 10),
		errChan:     make(chan error, 1),
	}
	go res.run()
	return res, nil
}

func (r *dataReader) Read(buf []byte) (int, error) {
	select {
	case data := <-r.dataChannel:
		if len(buf) < len(data) {
			return 0, errors.New("buf too small")
		}
		copy(buf, data)
		return len(data), nil
	case err := <-r.errChan:
		return 0, err
	}
}

func (r *dataReader) Close() error {
	r.closed.Store(true)
	return nil
}

func (r *dataReader) run() {
	buf := make([]byte, 2<<16)

	const headLen = binary.MaxVarintLen16
	var (
		readData   int
		dataLen    int
		shouldRead int
	)

	for {
		if r.closed.Load() {
			return
		}

		if dataLen == 0 {
			shouldRead = headLen - readData
		}
		r.conn.SetReadDeadline(time.Now().Add(time.Second))
		n, err := r.conn.Read(buf[readData:shouldRead])
		if err != nil {
			if errors.Is(err, os.ErrDeadlineExceeded) {
				continue
			}
			r.errChan <- err
			return
		}
		readData += n

		if readData == headLen {
			Len, i := binary.Varint(buf[:headLen])
			if i <= 0 {
				r.errChan <- errors.New("头部解析错误")
				return
			}
			dataLen = int(Len)
			shouldRead = dataLen + headLen
		} else if readData == shouldRead {
			sendData := make([]byte, dataLen)
			copy(sendData, buf[headLen:headLen+dataLen])
			r.dataChannel <- sendData
			readData = 0
		}
	}
}

// dataWriter 会将 Write 收到的数据打上长度头部并发送到 conn 中
type dataWriter struct {
	conn io.Writer
}

func NewDataWriter(conn io.Writer) (io.Writer, error) {
	if conn == nil {
		return nil, errors.New("conn is nil")
	}
	return &dataWriter{
		conn: conn,
	}, nil
}

func (w *dataWriter) Write(buf []byte) (int, error) {
	head := make([]byte, 3)
	binary.PutVarint(head, int64(len(buf)))
	_, err := w.conn.Write(head)
	if err != nil {
		return 0, err
	}
	return w.conn.Write(buf)
}
